#include "ros/ros.h"
#include "ros/time.h"
#include <iostream>
#include "tube_planning/mmpc.h"
#include "tube_planning/tube_planning.h"
// #include "tube_planning/polynomial.h"
#include "rapidjson/document.h"
#include <fstream>
#include "ros/package.h"
#include <string>
#include "tube_planning/config.h"
#include "OsqpEigen/OsqpEigen.h"

namespace Tube_planning
{
    Eigen::Array3d mmpc::init(ros::NodeHandle &nh, polyCurve traj)
    {
        traj_info = traj;
        readJson();
        model.x0 = Eigen::ArrayXd::Zero(model.nx);
        model.up = Eigen::ArrayXd::Zero(model.nu);
        mpc_var.x = Eigen::MatrixXd::Zero(model.nx, mpc_var.N+1);
        mpc_var.u = Eigen::MatrixXd::Zero(model.nu, mpc_var.N);
        mpc_var.xd = Eigen::MatrixXd::Zero(model.nx, mpc_var.N+1);
        mpc_var.ud = Eigen::MatrixXd::Zero(model.nu, mpc_var.N);
        mpc_var.Ak = Eigen::MatrixXd::Identity(model.nx,model.nx);
        mpc_var.Bk = Eigen::MatrixXd::Zero(model.nx, model.nu);
        begin_time =  ros::Time::now().toSec();
        pn_x = Polynomial(traj_info.Poly_x, traj_info.ts);
        pn_y = Polynomial(traj_info.Poly_y, traj_info.ts);
        pn_z = Polynomial(traj_info.Poly_z, traj_info.ts);

        Eigen::Array3d x0;
        init_time = ros::Time::now().toSec();
        delta_t = init_time - begin_time;
        x0(0) = pn_x.evaluatePolynomial(delta_t,0); 
        x0(1) = pn_y.evaluatePolynomial(delta_t,0); 
        x0(2) = pn_z.evaluatePolynomial(delta_t,0); 
        // std::cout << init_time << std::endl;
        return x0;
    }

    void mmpc::setBeginTime(const double reset_time)
    {
        begin_time = reset_time;
    }

    int mmpc::getRuntime()
    {
        double total_time = pn_x.getTotalTime();
        if (total_time > delta_t)
        {
            return 1;
        }
        else
        {
            return 0;
        }
    }

    void mmpc::getFeedback(const geometry_msgs::PoseStamped local_pos, 
                    const geometry_msgs::TwistStamped local_vel)
    {

        model.x0(0) = local_pos.pose.position.x;
        model.x0(1) = local_pos.pose.position.y;
        model.x0(2) = local_pos.pose.position.z;

        model.x0(3) = local_vel.twist.linear.x;
        model.x0(4) = local_vel.twist.linear.y;
        model.x0(5) = local_vel.twist.linear.z;
        // std::cout << "init state" << model.x0 << std::endl;

        model.up = mpc_var.u.block(0, 1, model.nu, 1);

    }

    Eigen::MatrixXd mmpc::getReference()
    {
        init_time = ros::Time::now().toSec();
        delta_t = init_time - begin_time;
        // desired state
        for (int i = 0; i < mpc_var.N + 1; i++)
        {
            double per_t = delta_t + mpc_var.Ts*i;
            // pos
            mpc_var.xd(0, i) = pn_x.evaluatePolynomial(per_t,0);
            mpc_var.xd(1, i) = pn_y.evaluatePolynomial(per_t,0);
            mpc_var.xd(2, i) = pn_z.evaluatePolynomial(per_t,0);
            // vel
            mpc_var.xd(3, i) = pn_x.evaluatePolynomial(per_t,1);
            mpc_var.xd(4, i) = pn_y.evaluatePolynomial(per_t,1);
            mpc_var.xd(5, i) = pn_z.evaluatePolynomial(per_t,1);

            if(i < mpc_var.N)
            {
                // input acc
                mpc_var.ud(0, i) = pn_x.evaluatePolynomial(per_t,2);
                mpc_var.ud(1, i) = pn_y.evaluatePolynomial(per_t,2);
                mpc_var.ud(2, i) = pn_z.evaluatePolynomial(per_t,2);
            }
        }
        // std::cout << "reference" << std::endl;
        // std::cout << mpc_var.xd << std::endl;
        return mpc_var.xd;

    }

    Eigen::ArrayXd mmpc::solver()
    {
        getMPCMaterix();
        quadprogInterface();
        Eigen::ArrayXd X = Eigen::ArrayXd::Zero(N*(model.nx + model.nu) + model.nx);
        for (int i = 0; i < N+1; i++)
        {
            X.segment(i*(model.nx + model.nu), model.nx) = mpc_var.x.block(0,i,model .nx,1);
            if (i < N)
            {
                X.segment(i*(model.nx + model.nu) + model.nx, model.nu) = mpc_var.u.block(0,i,model .nu,1);
            }
        }
        
        // X.segment(0, model.nx) = mpc_var.x.block(0, 1, model.nx, 1);
        // X.segment(model.nx, model.nu) = mpc_var.u.block(0, 0, model.nu, 1);
        // std::cout << "optimal" << std::endl;
        // std::cout << mpc_var.x << std::endl;
        // std::cout << mpc_var.xd << std::endl;
        // std::cout << mpc_var.x - mpc_var.xd << std::endl;
        return X;
    }

    void mmpc::getMPCMaterix()
    {
        int nxu = model.nu + model.nx;
        int nx = model.nx;
        int nu = model.nu;
        int i;
        for (i = 0; i < N; i++)
        {
            // cost for state
            stages_[i].Qk = Eigen::MatrixXd::Zero(nx,nx);
            stages_[i].Qk(0,0) = mpc_var.q_p;
            stages_[i].Qk(1,1) = mpc_var.q_p;
            stages_[i].Qk(2,2) = mpc_var.q_p;
            stages_[i].Qk(3,3) = mpc_var.q_v;
            stages_[i].Qk(4,4) = mpc_var.q_v;
            stages_[i].Qk(5,5) = mpc_var.q_v;
            // cost for input
            stages_[i].Rk = Eigen::MatrixXd::Zero(nu,nu);
            stages_[i].Rk(0,0) = mpc_var.q_a;
            stages_[i].Rk(1,1) = mpc_var.q_a;
            stages_[i].Rk(2,2) = mpc_var.q_a;
            // linear cost
            stages_[i].qfk = Eigen::ArrayXd::Zero(nx);
            stages_[i].qfk = -2 * stages_[i].Qk * mpc_var.xd.block(0, i, nx,1);
            stages_[i].rfk = Eigen::ArrayXd::Zero(nu);
            stages_[i].rfk = -2 * stages_[i].Rk * mpc_var.ud.block(0, i, nu,1);
            // model
            stages_[i].Ak = mpc_var.Ak;
            stages_[i].Ak(0,3) = mpc_var.Ts;
            stages_[i].Ak(1,4) = mpc_var.Ts;
            stages_[i].Ak(2,5) = mpc_var.Ts;
            stages_[i].Bk = mpc_var.Bk;
            stages_[i].Bk(3,0) = mpc_var.Ts;
            stages_[i].Bk(4,1) = mpc_var.Ts;
            stages_[i].Bk(5,2) = mpc_var.Ts;
            // bound
            stages_[i].lb = mpc_var.bound.block(0,0,nxu,1);
            stages_[i].ub = mpc_var.bound.block(0,1,nxu,1);
        }
        i = N;
        // cost for state
        stages_[i].Qk = Eigen::MatrixXd::Zero(nx,nx);
        stages_[i].Qk(0,0) = mpc_var.q_p;
        stages_[i].Qk(1,1) = mpc_var.q_p;
        stages_[i].Qk(2,2) = mpc_var.q_p;
        stages_[i].Qk(3,3) = mpc_var.q_v;
        stages_[i].Qk(4,4) = mpc_var.q_v;
        stages_[i].Qk(5,5) = mpc_var.q_v;
        // linear cost
        stages_[i].qfk = Eigen::ArrayXd::Zero(nx);
        stages_[i].qfk = -2 * stages_[i].Qk * mpc_var.xd.block(0, i, nx,1);
        // bound
        stages_[i].lb = mpc_var.bound.block(0,0,nx,1);
        stages_[i].ub = mpc_var.bound.block(0,1,nx,1);
    }

    void mmpc::quadprogInterface()
    {
        int nx = model.nx;
        int nu = model.nu;
        int nxu = nx + nu;
        // int nz = 2*nx + nu;
        Eigen::MatrixXd H = Eigen::MatrixXd::Zero(N*nxu + nx, N*nxu + nx);
        Eigen::VectorXd f = Eigen::VectorXd::Zero(N*nxu + nx);
        // Eigen::VectorXd Xd = Eigen::VectorXd::Zero(N*nxu + nx);
        for (int i = 0; i < N + 1; i++)
        {
            if (i < N)
            {
                H.block(i*nxu, i*nxu, nx, nx) = stages_[i].Qk;
                H.block(i*nxu + nx, i*nxu + nx, nu, nu) = stages_[i].Rk;
                f.segment(i*nxu, nx) = stages_[i].qfk;
                f.segment(i*nxu + nx, nu) = stages_[i].rfk;
                // Xd.segment(i*nxu, nx) = mpc_var.xd.block(0, i, nx, 1);
                // Xd.segment(i*nxu + nx, nu) = mpc_var.ud.block(0, i, nu, 1);
            }
            else if (i == N)
            {
                H.block(i*nxu, i*nxu, nx, nx) = stages_[i].Qk;
                f.segment(i*nxu, nx) = stages_[i].qfk;
                // Xd.segment(i*nxu, nx) = mpc_var.xd.block(0, i, nx, 1);
            }
        }

        H = 2 * H;
        // H =  H + H.transpose();
        // std::cout << "H" << H << std::endl;
        // std::cout << "Xd" << 0.5* Xd.transpose() * H * Xd + f.transpose() * Xd 
        // + 0.5 * Xd.transpose() * H * Xd << std::endl;
        int number_variables = N*nxu + nx;

        int number_init = nx + nu;
        Eigen::MatrixXd Aeq_1 = Eigen::MatrixXd::Identity(number_init, number_variables);
        Eigen::VectorXd lb_1(number_init);
        Eigen::VectorXd ub_1(number_init);
        lb_1 << model.x0, model.up;
        // lb.segment(N*nxu + nx + nx, nu) = model.up;
        ub_1 << model.x0, model.up;
        // ub.segment(N*nxu + nx + nx, nu) = model.up;

        int number_bound = N*nxu + nx;
        Eigen::MatrixXd Aeq_2 = Eigen::MatrixXd::Identity(number_bound, number_variables);
        Eigen::VectorXd lb_2 = Eigen::VectorXd::Zero(number_bound);
        Eigen::VectorXd ub_2 = Eigen::VectorXd::Zero(number_bound);
        for (int i = 0; i < mpc_var.N + 1; i++)
        {
            if (i < N)
            {
                lb_2.segment(i*nxu, nxu) = stages_[i].lb;
                ub_2.segment(i*nxu, nxu) = stages_[i].ub;
            }
            else if (i == N)
            {
                lb_2.segment(i*nxu, nx) = stages_[i].lb;
                ub_2.segment(i*nxu, nx) = stages_[i].ub;
            }
        }

        int number_model = mpc_var.N * nx;
        Eigen::MatrixXd Aeq_3 = Eigen::MatrixXd::Zero(number_model, number_variables);
        Eigen::VectorXd lb_3 = Eigen::VectorXd::Zero(number_model);
        Eigen::VectorXd ub_3 = Eigen::VectorXd::Zero(number_model);
        for (int i = 0; i < mpc_var.N; i++)
        {
            Aeq_3.block(i*nx, i*nxu, nx, nx) = stages_[i].Ak;
            Aeq_3.block(i*nx, i*nxu + nx, nx, nu) = stages_[i].Bk;
            Aeq_3.block(i*nx, i*nxu + nxu, nx, nx) = - Eigen::MatrixXd::Identity(nx, nx);
        }

        int number_cons = number_init + number_bound + number_model;

        Eigen::MatrixXd Aeq(number_cons, number_variables);
        Eigen::VectorXd lb(number_cons);
        Eigen::VectorXd ub(number_cons);
        Aeq << Aeq_1, Aeq_2, Aeq_3;
        lb << lb_1, lb_2, lb_3;
        ub << ub_1, ub_2, ub_3;


        Eigen::VectorXd QPSolution;
        Eigen::SparseMatrix<double> hessian;
        Eigen::SparseMatrix<double> linearMatrix;
        hessian = H.sparseView();
        linearMatrix = Aeq.sparseView();

        OsqpEigen::Solver solver;
        solver.settings()->setWarmStart(true);
        solver.settings()->setVerbosity(false);
        solver.data()->setNumberOfVariables(N*nxu + nx);
        solver.data()->setNumberOfConstraints(0);
        solver.data()->setNumberOfConstraints(number_cons);
        solver.data()->setHessianMatrix(hessian);
        solver.data()->setGradient(f);
        solver.data()->setLinearConstraintsMatrix(linearMatrix);
        solver.data()->setLowerBound(lb);
        solver.data()->setUpperBound(ub);


        solver.initSolver();
        solver.solveProblem();
        QPSolution = solver.getSolution();

        // std::cout << "Xo" << 0.5 * QPSolution.transpose() * H * QPSolution + f.transpose() * QPSolution 
        // + 0.5 * Xd.transpose() * H * Xd << std::endl;
        
        for (int i = 0; i < N + 1; i++)
        {
            mpc_var.x.block(0, i, nx, 1) = QPSolution.segment(i*nxu, nx);
            if (i < N)
            {
                mpc_var.u.block(0, i, nu, 1) = QPSolution.segment(i*nxu + nx, nu);
            }
        }
        
    }

    void mmpc::readJson()
    {
        std::string pkg_path = ros::package::getPath("tube_planning");
        std::string json_path = pkg_path + "/settings.json";
        std::ifstream file(json_path,std::ios_base::in);
        if (!file.is_open())
        {
            std::cout << "can not open json failed." << std::endl;
        }
        std::string json_content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
        file.close();

        rapidjson::Document dom;
	    if (dom.Parse(json_content.c_str()).HasParseError())
        {
            std::cout << "parse error" << std::endl;
        }
        if (dom.HasMember("MPC"))
        {
            mpc_var.N = dom["MPC"]["N"].GetInt();
            // std::cout << "mpc_var.N:" << mpc_var.N << std::endl;
            mpc_var.Ts = dom["MPC"]["Ts"].GetDouble();
            // std::cout << "mpc_var.Ts:" << mpc_var.Ts << std::endl;
            const rapidjson::Value& array_lb = dom["MPC"]["bound"]["lb"].GetArray();
            const rapidjson::Value& array_ub = dom["MPC"]["bound"]["ub"].GetArray();
            int arraySize = array_lb.Size();
            mpc_var.bound.resize(arraySize,2);
            for(int i = 0; i < arraySize; i ++)
            {
                mpc_var.bound(i,0) = array_lb[i].GetDouble();
                mpc_var.bound(i,1) = array_ub[i].GetDouble();
            }
            // std::cout << "mpc_var.bound: " << mpc_var.bound.transpose() << std::endl;
            mpc_var.q_p = dom["MPC"]["qp"].GetDouble();
            mpc_var.q_v = dom["MPC"]["qv"].GetDouble();
            mpc_var.q_a = dom["MPC"]["qa"].GetDouble();
            // std::cout << "mpc_var.qp: " << mpc_var.q_p << std::endl;
            // std::cout << "mpc_var.qv: " << mpc_var.q_v << std::endl;
            // std::cout << "mpc_var.qa: " << mpc_var.q_a << std::endl;
            model.nx = dom["Model"]["nx"].GetInt();
            // std::cout << model.nx << std::endl;
            model.nu = dom["Model"]["nu"].GetInt();
            model.raduis = dom["Model"]["radius"].GetDouble();
        }
    }
}